import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from heapq import heappush, heappop
from tqdm import tqdm

# ============================================
# 参数定义（改进楼梯高度和移动逻辑）
# ============================================
FLOORS = 10          # 减少楼层数方便观察
GRID_SIZE = 15       # 增大网格尺寸
STAIR_WIDTH = 2      # 楼梯宽度
STAIR_DIRECTION = (1, 1)  # 楼梯斜向方向 (dx, dy)
STAIR_HEIGHT = 3     # 每级台阶高度系数（论文中的斜面距离）
PEOPLE_PER_FLOOR = 5
V0_MEAN = 1.0
V0_STD = 0.1
a = 0.71
b = 0.51

# ============================================
# 三维静态场计算（改进楼梯斜向移动）
# ============================================
def calculate_static_field_3d():
    field = np.full((GRID_SIZE, GRID_SIZE, FLOORS), np.inf)
    exit_pos = (GRID_SIZE//2, GRID_SIZE//2, 0)
    field[exit_pos] = 0
    
    heap = []
    heappush(heap, (0, *exit_pos))
    
    # 楼梯移动方向集（斜向）
    stair_directions = []
    for dz in [-1, 0]:
        dx, dy = STAIR_DIRECTION
        if dz == -1:  # 仅允许下楼时斜向移动
            stair_directions.extend([
                (dx, dy, dz), (-dx, dy, dz),
                (dx, -dy, dz), (-dx, -dy, dz)
            ])
    
    while heap:
        dist, x, y, z = heappop(heap)
        # 常规移动方向（平面）
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                if dx == 0 and dy == 0: continue
                nx, ny, nz = x+dx, y+dy, z
                if 0 <= nx < GRID_SIZE and 0 <= ny < GRID_SIZE:
                    new_dist = dist + np.sqrt(dx**2 + dy**2)
                    if new_dist < field[nx, ny, nz]:
                        field[nx, ny, nz] = new_dist
                        heappush(heap, (new_dist, nx, ny, nz))
        
        # 楼梯斜向移动（三维）
        for dx, dy, dz in stair_directions:
            nx, ny, nz = x+dx, y+dy, z+dz
            if 0 <= nx < GRID_SIZE and 0 <= ny < GRID_SIZE and nz >= 0:
                # 计算斜面距离（考虑台阶高度）
                step_length = np.sqrt(dx**2 + dy**2 + (dz*STAIR_HEIGHT)**2)
                new_dist = dist + step_length
                if new_dist < field[nx, ny, nz]:
                    field[nx, ny, nz] = new_dist
                    heappush(heap, (new_dist, nx, ny, nz))
    
    return field

# ============================================
# 改进的元胞类（考虑斜向移动）
# ============================================
class Person:
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z
        self.v0 = np.clip(np.random.normal(V0_MEAN, V0_STD), 0.5, 1.5)
        self.total_distance = 0.0
        self.step_counter = 0
        self.time_step = max(1, int(1 / self.v0))
    
    def update_speed(self):
        c = 1 if self.z > 0 else 2.7
        exponent = a * (self.total_distance ** b)
        speed_factor = 1 - 0.01 * c * np.exp(exponent) + 0.01
        new_v = self.v0 * speed_factor
        self.time_step = max(1, int(1 / new_v))  # 保证最小移动间隔
    
    def move(self, static_field):
        min_dist = np.inf
        best_dir = (0, 0, 0)
        
        # 平面移动方向
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                if dx == 0 and dy == 0: continue
                nx, ny, nz = self.x+dx, self.y+dy, self.z
                if 0 <= nx < GRID_SIZE and 0 <= ny < GRID_SIZE:
                    if static_field[nx, ny, nz] < min_dist:
                        min_dist = static_field[nx, ny, nz]
                        best_dir = (dx, dy, 0)
        
        # 楼梯斜向移动
        for dx, dy, dz in [
            (STAIR_DIRECTION[0], STAIR_DIRECTION[1], -1),
            (-STAIR_DIRECTION[0], STAIR_DIRECTION[1], -1),
            (STAIR_DIRECTION[0], -STAIR_DIRECTION[1], -1),
            (-STAIR_DIRECTION[0], -STAIR_DIRECTION[1], -1)
        ]:
            nx, ny, nz = self.x+dx, self.y+dy, self.z+dz
            if 0 <= nx < GRID_SIZE and 0 <= ny < GRID_SIZE and nz >= 0:
                if static_field[nx, ny, nz] < min_dist:
                    min_dist = static_field[nx, ny, nz]
                    best_dir = (dx, dy, dz)
        
        # 执行移动
        self.x += best_dir[0]
        self.y += best_dir[1]
        self.z += best_dir[2]
        self.total_distance += np.sqrt(
            best_dir[0]**2 + best_dir[1]**2 + (abs(best_dir[2])*STAIR_HEIGHT**2)
        )
        self.update_speed()

# ============================================
# 初始化模型
# ============================================
static_field = calculate_static_field_3d()
people = []
for z in range(FLOORS):
    for _ in range(PEOPLE_PER_FLOOR):
        while True:
            x = np.random.randint(2, GRID_SIZE-2)
            y = np.random.randint(2, GRID_SIZE-2)
            # 确保初始不在楼梯附近
            if (abs(x - GRID_SIZE//2) > STAIR_WIDTH or 
                abs(y - GRID_SIZE//2) > STAIR_WIDTH):
                people.append(Person(x, y, z))
                break

# ============================================
# 改进可视化（显示楼梯斜向结构）
# ============================================
def visualize(people, step):
    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制斜向楼梯结构
    for z in range(FLOORS):
        x_start = GRID_SIZE//2 - STAIR_WIDTH//2
        y_start = GRID_SIZE//2 - STAIR_WIDTH//2
        for i in range(STAIR_WIDTH):
            for j in range(STAIR_WIDTH):
                # 绘制斜线路径
                ax.plot(
                    [x_start+i, x_start+i + STAIR_DIRECTION[0]],
                    [y_start+j, y_start+j + STAIR_DIRECTION[1]],
                    [z, max(z-1, 0)],
                    c='gray', alpha=0.3
                )
    
    # 绘制人员
    xs, ys, zs = [], [], []
    for p in people:
        if p.z >= 0 and (p.x != GRID_SIZE//2 or p.y != GRID_SIZE//2 or p.z != 0):
            xs.append(p.x)
            ys.append(p.y)
            zs.append(p.z)
    ax.scatter(xs, ys, zs, c='red', s=50, depthshade=False)
    
    ax.set_xlim(0, GRID_SIZE)
    ax.set_ylim(0, GRID_SIZE)
    ax.set_zlim(0, FLOORS)
    ax.set_box_aspect([1,1,1])
    ax.set_title(f'Step: {step}  Remaining: {len(xs)}')
    plt.show()

# ============================================
# 主循环（改进时间控制）
# ============================================
max_steps = 500
for step in tqdm(range(max_steps)):
    # 仅更新可移动人员
    active_people = [p for p in people if not (p.z == 0 and p.x == GRID_SIZE//2 and p.y == GRID_SIZE//2)]
    
    for p in active_people:
        p.step_counter += 1
        if p.step_counter >= p.time_step:
            p.move(static_field)
            p.step_counter = 0
    
    # 实时可视化
    if step % 10 == 0:  # 提高可视化频率
        visualize(people, step)
    
    if len(active_people) == 0:
        print(f"All evacuated at step {step}")
        break

visualize(people, step)  # 最终状态